#!/usr/bin/env python

import roslib; roslib.load_manifest('team4')
import rospy
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import arm_planner
from arm_mover import ArmMover

from pr2_controllers_msgs.msg import JointTrajectoryAction
from geometry_msgs.msg import PointStamped, PoseStamped
from trajectory_msgs.msg import JointTrajectoryPoint, JointTrajectory
import numpy as np
from tf.transformations import *

class BasketManip():
   def __init__(self):
      self.basket_location = PointStamped()
      self.last_grasps = None
      self.basket = None
      self.l_arm_mover = ArmMover('left_arm')
      self.r_arm_mover = ArmMover('right_arm')

   def find_basket(self):
      # Go to predetermined basket location?
      det = TablewareDetection()
      head_point = PointStamped()
      head_point.header.frame_id = '/torso_lift_link'
      head_point.point.x = 0.6
      res = det.detect_objects(point_head_at=head_point)
      
      for pg in res.pickup_goals:
         if pg.label == 'graspable':
            # What happens with more than one graspable? (i.e. stack?)
            self.basket = pg
      return

      
   def getArmIK(self, pose=PoseStamped(), arm='right_arm'):
      arm_planner=[]
      if arm == 'right_arm':
	      arm_planner = self.r_arm_mover._arm_planner
      elif arm == 'left_arm':
	      arm_planner = self.l_arm_mover._arm_planner
      print "arm_planner " + str(arm_planner)
      ik_sol = arm_planner.get_ik(pose)
      return ik_sol
 
   def empty_basket(self):
      rospy.loginfo("Tip over basket not implemented")
      
   
   def pick_up_basket(self):
      rospy.loginfo("Pick up basket not implemented")
      

   def put_down_object(self):
      rospy.loginfo("Put down basket not implemented")

   def place_object_in_basket(self, obj):
      rospy.loginfo("Place object in basket not implemented") 
   


if __name__ == "__main__":
   rospy.init_node("basket_controller")
   rospy.loginfo("basket")
   b = BasketManip()
   
   j = JointTrajectory()
   
   right_arm_joints = ['r_shoulder_pan_joint','r_shoulder_lift_joint','r_upper_arm_roll_joint','r_elbow_flex_joint','r_forearm_roll_joint','r_wrist_flex_joint','r_wrist_roll_joint']
   left_arm_joints = ['l_shoulder_pan_joint','l_shoulder_lift_joint','l_upper_arm_roll_joint','l_elbow_flex_joint','l_forearm_roll_joint','l_wrist_flex_joint','l_wrist_roll_joint' ]

   j.joint_names = right_arm_joints + left_arm_joints
   offset=1.0
   for i in [0.1, 0.0, -0.1]:
      p = PoseStamped()
      p.header.frame_id = '/torso_lift_link'
      p.pose.position.x = 0.5
      p.pose.position.y = -0.4
      p.pose.position.z = i
      quat=quaternion_from_euler(0,np.pi/2.0, 0)
      p.pose.orientation.x= quat[0]
      p.pose.orientation.y =quat[1]
      p.pose.orientation.z =quat[2]
      p.pose.orientation.w =quat[3]

      soln_r = b.getArmIK(p,'right_arm')
      #rospy.loginfo("Arm IK solutions for right_arm: {0}".format(soln_r))
      solution = []
      for k in right_arm_joints:
         solution.append(soln_r.solution.joint_state.position[soln_r.solution.joint_state.name.index(k)])

      p.pose.position.y = 0.4
      soln_l = b.getArmIK(p, 'left_arm')
      for k in left_arm_joints:
         solution.append(soln_l.solution.joint_state.position[soln_l.solution.joint_state.name.index(k)])
#      rospy.loginfo("Arm IK solutions for left_arm: {0}".format(soln_l))

   
      tp = JointTrajectoryPoint()
      tp.time_from_start = rospy.Duration(3.0+offset)
      offset=offset+offset
      tp.positions = solution
      print soln_r
      print soln_l
      tp.velocities = list(np.zeros(len(tp.positions)))
      j.points.append(tp)

   j.header.stamp=rospy.Time.now()+rospy.Duration(1.0)
   b.l_arm_mover.execute_two_arm_trajectory(j)
    
   
     
